# utils/file_parser.py
# 全能文件解析器，可以将多种格式的文件统一转换为纯文本

import os
import json
import fitz  # PyMuPDF别名
import docx  # python-docx别名
from PIL import Image
import pytesseract

class FileParser:
    """
    一个通用的文件解析器，支持 pdf, docx, txt, json, png, jpg 等格式。
    所有文件最终都被转换为纯文本字符串。
    """
    def __init__(self, tesseract_lang: str = 'chi_sim+eng'):
        """
        初始化解析器。
        :param tesseract_lang: 为OCR指定的语言, 例如 'eng' (英文), 'chi_sim' (简体中文)。
                              可以组合，如 'chi_sim+eng'。
        """
        self.tesseract_lang = tesseract_lang

    def extract_text(self, file_path: str) -> str:
        """
        主分发方法。根据文件扩展名，调用相应的处理函数。
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件未找到: {file_path}")

        # 获取文件扩展名，并转换为小写以统一处理
        _, extension = os.path.splitext(file_path)
        ext_lower = extension.lower()

        try:
            if ext_lower == '.pdf':
                return self._extract_from_pdf(file_path)
            elif ext_lower == '.docx':
                return self._extract_from_docx(file_path)
            elif ext_lower in ['.png', '.jpg', '.jpeg', '.tiff', '.bmp']:
                return self._extract_from_image(file_path)
            elif ext_lower == '.txt':
                return self._extract_from_txt(file_path)
            elif ext_lower == '.json':
                return self._extract_from_json(file_path)
            else:
                raise ValueError(f"不支持的文件类型: {ext_lower}")
        except Exception as e:
            # 捕获所有处理过程中的异常，并附加文件名信息
            raise RuntimeError(f"处理文件 '{os.path.basename(file_path)}' 时出错: {e}")

    def _extract_from_pdf(self, file_path: str) -> str:
        """从PDF文件中提取文本。"""
        with fitz.open(file_path) as doc:
            all_text = [page.get_text().strip() for page in doc]
            return "\n\n".join(all_text)

    def _extract_from_docx(self, file_path: str) -> str:
        """从DOCX文件中提取文本。"""
        doc = docx.Document(file_path)
        all_text = [p.text for p in doc.paragraphs]
        return "\n".join(all_text)

    def _extract_from_image(self, file_path: str) -> str:
        """从图像文件中使用OCR提取文本。"""
        # 需要系统已正确安装 Tesseract-OCR 引擎
        try:
            return pytesseract.image_to_string(Image.open(file_path), lang=self.tesseract_lang)
        except pytesseract.TesseractNotFoundError:
            raise RuntimeError(
                "Tesseract-OCR 未安装或未在系统PATH中。请访问 https://github.com/tesseract-ocr/tesseract 进行安装。"
            )

    def _extract_from_txt(self, file_path: str) -> str:
        """从纯文本文件中提取文本。"""
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()

    def _extract_from_json(self, file_path: str) -> str:
        """
        从JSON文件中提取内容并转换为格式化的纯文本。
        这有助于LLM理解其结构。
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        # 将JSON数据美化打印为字符串，保留缩进和结构
        return json.dumps(data, ensure_ascii=False, indent=2)